package com.em;

import be.ac.ulg.montefiore.run.distributions.GaussianDistribution;
import be.ac.ulg.montefiore.run.jahmm.Observation;
import be.ac.ulg.montefiore.run.jahmm.ObservationReal;
import be.ac.ulg.montefiore.run.jahmm.OpdfGaussian;
import be.ac.ulg.montefiore.run.jahmm.apps.cli.CommandLineArguments;
import be.ac.ulg.montefiore.run.jahmm.apps.cli.Types;
import be.ac.ulg.montefiore.run.jahmm.apps.cli.WrongArgumentsException;
import be.ac.ulg.montefiore.run.jahmm.io.FileFormatException;
import com.utils.Constants;
import com.utils.SequenceIterator;

import java.util.*;

/**
 * This class is used to calculate, using EM algorithm the gaussian mixture given specific input sequence
 * Created with IntelliJ IDEA.
 * User: bigmar
 * Date: 12/4/14
 * Time: 3:16 PM
 * To change this template use File | Settings | File Templates.
 */
public class EMForGaussians
{
    private int numberOfModels;
    private int numberOfObservations;
    private SequenceIterator sequenceIterator;
    public EMForGaussians(int numberOfModels, int numberOfObservations, SequenceIterator sequenceIterator)
    {
        this.numberOfModels=numberOfModels;
        this.numberOfObservations=numberOfObservations;
        this.sequenceIterator=sequenceIterator;

    }

    /**
     * We run a maximum of 5 valid iterations, but just to make sure, we also allow a total of 1000 iterations (In case all of them fail
     * @return
     */
    public EmForGaussiansResult run()
    {
        int maxIterations=1000;
        EmForGaussiansResult retVal=null;

        int validIterations=0;
        double logLikelihood= Integer.MIN_VALUE;
        for (int i=0; i<maxIterations&&validIterations<Constants.numberOfRandomSelectedInitValuesForEm; ++i)
        {
            EmForGaussiansResult tmp=runInner(logLikelihood);
            if (tmp!=null)
            {
                validIterations++;
                if (retVal==null||tmp.getLogLikelihood()>retVal.getLogLikelihood())
                    retVal=tmp;

                logLikelihood=retVal.getLogLikelihood();
            }
            System.out.println("iteration number "+i);


        }
        return retVal;
    }
    private EmForGaussiansResult runInner(double prevRunLogLikelihood)
    {


        double[] mixtureWeights=initRandomMixtureWeights(numberOfModels);
        List<OpdfGaussian> gaussianList=initRandomGaussianDistributions(mixtureWeights);


        double[][] weights=new double[numberOfObservations][numberOfModels];
        double prevLogLikelihood=Integer.MAX_VALUE;
        for (int i=0; i< Constants.numberOfEmIterations; ++i)
        {
            calcExpectation(mixtureWeights, gaussianList, weights);
            if (!calcMaximization(mixtureWeights,gaussianList, weights))
                return null;
            double logLikelihood=calcLikelihood(mixtureWeights, gaussianList);

            if (Math.abs(prevLogLikelihood-logLikelihood)<Constants.minLogLikelihood)
            {
                System.out.println(""+i+")"+logLikelihood);
                prevLogLikelihood=logLikelihood;

                break;
            }
            if (Math.abs(prevLogLikelihood-logLikelihood)<0.01&&1.1*prevRunLogLikelihood>logLikelihood)
            {
                prevLogLikelihood=logLikelihood;

                break;
            }

            prevLogLikelihood=logLikelihood;
        }

        return new EmForGaussiansResult(mixtureWeights, gaussianList, prevLogLikelihood);

    }
    private double calcLikelihood(double[] mixtureWeights, List<OpdfGaussian> gaussianList)
    {
        int i=0;
        double logLikelihood=0;
        sequenceIterator.init();
        while(sequenceIterator.hasNext())
        {
            ObservationReal value=((ObservationReal) sequenceIterator.next());
            ++i;
            double oneSeqLogLikelihood=0;

            for (int j=0; j<numberOfModels; ++j)
            {

                oneSeqLogLikelihood+=  mixtureWeights[j]*gaussianList.get(j).probability(value);
            }
            oneSeqLogLikelihood=Math.log(oneSeqLogLikelihood);
            logLikelihood+=oneSeqLogLikelihood;
        }

        return logLikelihood;
    }

    /**
     * Calculate maximization
     * @param mixtureWeights
     * @param gaussianList
     * @param weights
     */
    private boolean calcMaximization(double[] mixtureWeights, List<OpdfGaussian> gaussianList, double[][] weights)
    {

        for (int modelIndex=0; modelIndex<numberOfModels; ++modelIndex)
        {

            // Calculate new mixture weights
            mixtureWeights[modelIndex]=0;
            for (int observationIndex=0; observationIndex<numberOfObservations;++observationIndex)
            {
                mixtureWeights[modelIndex]+=weights[observationIndex][modelIndex];
            }
            mixtureWeights[modelIndex]/=numberOfObservations;
            if (mixtureWeights[modelIndex]==0)
                return false;

            // Calculate new expectations
            double newGaussianExpectation=0;
            sequenceIterator.init();
            int i=0;
            while(sequenceIterator.hasNext())
            {
                double value=((ObservationReal) sequenceIterator.next()).value;
                ++i;
                newGaussianExpectation+=weights[sequenceIterator.getCurrentObservation()][modelIndex]*value;
            }
            newGaussianExpectation/=mixtureWeights[modelIndex];
            newGaussianExpectation/=numberOfObservations;

            // Calculate new variance
            double newGaussianVariance=0;
            sequenceIterator.init();
            while(sequenceIterator.hasNext())
            {
                double value=((ObservationReal) sequenceIterator.next()).value;

                newGaussianVariance+=weights[sequenceIterator.getCurrentObservation()][modelIndex]*(value - newGaussianExpectation)*(value - newGaussianExpectation);

            }
            newGaussianVariance/=mixtureWeights[modelIndex];
            newGaussianVariance/=numberOfObservations;

            if (newGaussianVariance<=0)
            {
                return false;
            }
            gaussianList.set(modelIndex, new OpdfGaussian(newGaussianExpectation, newGaussianVariance));
        }
        return true;
    }

    /**
     * The expectation step in the EM algorithm
     * @param mixtureWeights
     * @param gaussianList
     */
    private double[][] calcExpectation(double[] mixtureWeights, List<OpdfGaussian> gaussianList,double[][] weights)
    {

        // Index for all sequences
        int observationIndex=0;
        sequenceIterator.init();
        while(sequenceIterator.hasNext())
        {
            ObservationReal value=((ObservationReal) sequenceIterator.next());


            double sumForNormalization=0;
            for (int modelIndex=0; modelIndex<numberOfModels; ++modelIndex)
            {
                weights[observationIndex][modelIndex]=mixtureWeights[modelIndex]*
                        gaussianList.get(modelIndex).probability(value);
                if (Double.isNaN(weights[observationIndex][modelIndex]))
                    weights[observationIndex][modelIndex]=mixtureWeights[modelIndex]*
                            gaussianList.get(modelIndex).probability(value);
                sumForNormalization+=weights[observationIndex][modelIndex];
            }
            for (int modelIndex=0; modelIndex<numberOfModels; ++modelIndex)
            {
                if (sumForNormalization!=0)
                    weights[observationIndex][modelIndex]/=sumForNormalization;
            }

            ++observationIndex;
        }
        return weights;
    }
    private List<OpdfGaussian> initRandomGaussianDistributions(double[] mixtureWeights)
    {
        // Initiate expectation
        double expectation=0;
        sequenceIterator.init();
        SortedSet<Double> sequence=new TreeSet<Double>();

        while(sequenceIterator.hasNext())
        {
            double value=((ObservationReal) sequenceIterator.next()).value;
            sequence.add(value);
        }
        expectation/=numberOfObservations;

        Iterator<Double> iterator=sequence.iterator();
        List<OpdfGaussian> gaussianList=new ArrayList<OpdfGaussian>(numberOfModels);

        int totalNumberOfObjectsSoFar=0;
        for (int i=0; i<mixtureWeights.length; ++i)
        {
            int numberOfObjects;
            if (i==mixtureWeights.length-1)
                numberOfObjects=sequence.size()-totalNumberOfObjectsSoFar;
            else
                numberOfObjects=(int)(mixtureWeights[i]*sequence.size());

            totalNumberOfObjectsSoFar+=numberOfObjects;

            for (int j=0; j<numberOfObjects; ++j)
            {
                double value=iterator.next();

                expectation+=value;
            }
            expectation/=numberOfObjects;
            gaussianList.add(new OpdfGaussian(expectation, 2*Math.random()));

        }
        return gaussianList;
    }
    /**
     * Initialize random mixture weights
     * Array of positive numbers summing up to 1
     * @param numberOfModels
     * @return
     */
    private double[] initRandomMixtureWeights(int numberOfModels)
    {
        double[] mixtureWeights=new double[numberOfModels];
        double sum=0;
        for (int i=0; i<numberOfModels; ++i)
        {
            mixtureWeights[i]=Math.random();
            sum+=mixtureWeights[i];
        }
        for (int i=0; i<numberOfModels; ++i)
        {
            mixtureWeights[i]/=sum;
        }
        return mixtureWeights;
    }
}
